import os
import sys
import argparse
import requests
import pandas as pd
from tqdm import tqdm
import time

def download_ximalaya_audios(excel_path, output_folder, delay=1):
    """
    从Excel下载喜马拉雅音频文件
    
    :param excel_path: Excel文件路径
    :param output_folder: 音频保存目录
    :param delay: 每次下载间隔(秒)，避免请求过快
    """
    # 创建输出目录
    os.makedirs(output_folder, exist_ok=True)
    
    try:
        # 读取Excel文件（跳过第一行表头）
        df = pd.read_excel(excel_path, header=0)
    except Exception as e:
        print(f"读取Excel文件失败: {e}")
        return False

    # 验证列名
    if 'file_name' not in df.columns or 'URL' not in df.columns:
        print("Excel必须包含'file_name'和'URL'列")
        return False

    # 喜马拉雅专用Headers
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
        'Referer': 'https://www.ximalaya.com/'
    }

    success = 0
    fail = 0
    
    for index, row in tqdm(df.iterrows(), total=len(df), desc="下载进度"):
        filename = str(row['file_name']).strip()
        url = str(row['URL']).strip()
        
        if not filename or not url or pd.isna(url):
            print(f"第{index+2}行数据不完整，跳过")
            fail += 1
            continue
        
        # 确保文件名有.m4a扩展名（喜马拉雅音频格式）
        if not filename.lower().endswith('.m4a'):
            filename += '.m4a'
        
        try:
            filepath = os.path.join(output_folder, filename)
            
            # 检查文件是否已存在
            if os.path.exists(filepath):
                print(f"文件已存在，跳过: {filename}")
                success += 1
                continue
                
            # 发送带Headers的请求
            response = requests.get(
                url,
                headers=headers,
                stream=True,
                timeout=30
            )
            response.raise_for_status()
            
            # 流式下载大文件
            with open(filepath, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            
            success += 1
            time.sleep(delay)  # 避免请求过快
            
        except requests.exceptions.RequestException as e:
            print(f"下载失败【{filename}】: {str(e)}")
            fail += 1
        except Exception as e:
            print(f"处理【{filename}】时出错: {str(e)}")
            fail += 1
    
    print(f"\n下载完成！成功: {success}个, 失败: {fail}个")
    return True

def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='喜马拉雅音频下载工具')
    parser.add_argument('-i', '--input', required=True, help='输入的Excel文件路径')
    parser.add_argument('-o', '--output', default='downloads', help='音频保存目录（默认为当前目录下的downloads文件夹）')
    parser.add_argument('-d', '--delay', type=float, default=1.0, help='下载间隔时间（秒），默认1秒')
    parser.add_argument('-r', '--retry', type=int, default=3, help='失败重试次数，默认3次')
    
    args = parser.parse_args()
    
    # 添加重试机制
    for attempt in range(args.retry):
        try:
            print(f"\n开始下载（第{attempt+1}次尝试）...")
            if download_ximalaya_audios(args.input, args.output, args.delay):
                break
        except Exception as e:
            if attempt == args.retry - 1:
                print(f"多次尝试后仍失败: {e}")
                sys.exit(1)
            else:
                print(f"第{attempt+1}次尝试失败，{5*(attempt+1)}秒后重试...")
                time.sleep(5 * (attempt + 1))

if __name__ == "__main__":
    main()
